import torch
import torch.utils.data
from src.const import base_path
import numpy as np
import cv2
from torchvision import transforms
import matplotlib.pyplot as plt
import pandas as pd
from skimage import io, transform


class Rescale(object):
    """Rescale the image in a sample to a given size.

    Args:
        output_size (tuple or int): Desired output size. If tuple, output is
            matched to output_size. If int, smaller of image edges is matched
            to output_size keeping aspect ratio the same.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        self.output_size = output_size

    def __call__(self, image, landmarks):
        h, w = image.shape[:2]
        if isinstance(self.output_size, int):
            if h > w:
                new_h, new_w = self.output_size * h / w, self.output_size
            else:
                new_h, new_w = self.output_size, self.output_size * w / h
        else:
            new_h, new_w = self.output_size

        new_h, new_w = int(new_h), int(new_w)

        img = transform.resize(image, (new_h, new_w), mode='constant')

        # h and w are swapped for landmarks because for images,
        # x and y axes are axis 1 and 0 respectively
        landmarks = landmarks * [new_w / w, new_h / h]

        return img, landmarks


class RandomCrop(object):
    """Crop randomly the image in a sample.

    Args:
        output_size (tuple or int): Desired output size. If int, square crop
            is made.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        if isinstance(output_size, int):
            self.output_size = (output_size, output_size)
        else:
            assert len(output_size) == 2
            self.output_size = output_size

    def __call__(self, image, landmarks):
        h, w = image.shape[:2]
        new_h, new_w = self.output_size

        top = np.random.randint(0, h - new_h)
        left = np.random.randint(0, w - new_w)

        image = image[top: top + new_h,
                      left: left + new_w]

        landmarks = landmarks - [left, top]

        return image, landmarks

class CenterCrop(object):

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        if isinstance(output_size, int):
            self.output_size = (output_size, output_size)
        else:
            assert len(output_size) == 2
            self.output_size = output_size

    def __call__(self, image, landmarks):
        h, w = image.shape[:2]
        new_h, new_w = self.output_size

        top = int((h - new_h) / 2)
        left = int((w - new_w) / 2)

        image = image[top: top + new_h,
                      left: left + new_w]

        landmarks = landmarks - [left, top]

        return image, landmarks


class RandomFlip(object):

    def __call__(self, image, landmarks):
        h, w = image.shape[:2]
        if np.random.rand() > 0.5:
            image = np.fliplr(image)
            landmarks[:, 0] = w - landmarks[:, 0]

        return image, landmarks


class CheckLandmarks(object):

    def __call__(self, image, landmark_vis, landmark_pos):
        h, w = image.shape[:2]
        landmark_vis = landmark_vis.copy()
        landmark_pos = landmark_pos.copy()
        for i, vis in enumerate(landmark_vis):
            if (landmark_pos[i, 0] < 0) or (landmark_pos[i, 0] >= w) or (landmark_pos[i, 1] < 0) or (landmark_pos[i, 1] >= h):
                landmark_vis[i] = 0  # 不可见
        for i, vis in enumerate(landmark_vis):
            if vis == 0:
                landmark_pos[i, :] = 0
        return landmark_vis, landmark_pos


class LandmarksNormalize(object):

    def __call__(self, image, landmark_pos):
        h, w = image.shape[:2]
        landmark_pos = landmark_pos / [float(w), float(h)]
        return landmark_pos

class LandmarksUnNormalize(object):

    def __call__(self, image, landmark_pos):
        h, w = image.shape[:2]
        landmark_pos = landmark_pos * [float(w), float(h)]
        return landmark_pos

class DeepFashionCAPDataset(torch.utils.data.Dataset):
    '''
    mode: 'RANDOM', 'CENTER', 'NOCHANGE'
    '''

    def __init__(self, df, mode, base_path=base_path):
        self.df = df
        self.base_path = base_path
        self.to_tensor = transforms.ToTensor()  # pytorch使用c x h x w的格式转换
        self.rescale = Rescale(256)
        self.rescale_largest_center = Rescale(224)
        self.center_crop = CenterCrop(224)
        self.random_crop = RandomCrop(224)
        # self.random_flip = RandomFlip()
        self.check_landmarks = CheckLandmarks()
        self.to_tensor = transforms.ToTensor()
        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                              std=[0.229, 0.224, 0.225])
        self.landmarks_normalize = LandmarksNormalize()
        self.landmarks_unnormalize = LandmarksUnNormalize()
        self.mode = mode
        assert self.mode in ['RANDOM', 'CENTER', 'LARGESTCENTER']

        # for vis
        self.unnormalize = transforms.Normalize(mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225],
                                                std=[1 / 0.229, 1 / 0.224, 1 / 0.225])
        self.to_pil = transforms.ToPILImage()

    def plot_sample(self, i):
        sample = self[i]
        if isinstance(sample['image'], torch.Tensor):
            image = self.unnormalize(sample['image'])
            image = self.to_pil(image.float())
            image = np.array(image)
        plt.figure(dpi=200)
        plt.imshow(image)
        for i, vis in enumerate(sample['landmark_vis']):
            if (vis == 1):
                plt.scatter([sample['landmark_pos'][i, 0]], [sample['landmark_pos'][i, 1]], s=20, marker='.', c='g')
            else:
                plt.scatter([sample['landmark_pos'][i, 0]], [sample['landmark_pos'][i, 1]], s=20, marker='x', c='r')

    def __getitem__(self, i):
        sample = self.df.iloc[i]
        image = io.imread(base_path + sample['image_name'])
        category_label = sample['category_label']  # 0开始，这是一个1-of-K的单分类问题
        landmark_vis = sample.filter(regex='lm.*vis').astype(np.int64).values  # 1表示可见，0表示不可见
        landmark_pos_x = sample.filter(regex='lm.*x').astype(np.int64).values.reshape(-1, 1)
        landmark_pos_y = sample.filter(regex='lm.*y').astype(np.int64).values.reshape(-1, 1)
        landmark_pos = np.concatenate([landmark_pos_x, landmark_pos_y], axis=1)
        attr = sample.filter(regex='attr.*').astype(np.int64).values  # 0表示没有，1表示有
        category_type = sample['category_type']

        if self.mode == 'RANDOM':
            image, landmark_pos = self.rescale(image, landmark_pos)
            image, landmark_pos = self.random_crop(image, landmark_pos)
            # image, landmark_pos = self.random_flip(image, landmark_pos)
        elif self.mode == 'CENTER':
            image, landmark_pos = self.rescale(image, landmark_pos)
            image, landmark_pos = self.center_crop(image, landmark_pos)
        elif self.mode == 'LARGESTCENTER':
            image, landmark_pos = self.rescale_largest_center(image, landmark_pos)
            image, landmark_pos = self.center_crop(image, landmark_pos)
        else:
            raise NotImplementedError
        landmark_vis, landmark_pos = self.check_landmarks(image, landmark_vis, landmark_pos)

        landmark_pos = landmark_pos.astype(np.float32)
        landmark_pos_normalized = self.landmarks_normalize(image, landmark_pos).astype(np.float32)

        image = image.copy()  # 部分numpy存储模式不支持

        image = self.to_tensor(image)
        image = self.normalize(image)
        image = image.float()

        ret = {}
        ret['image'] = image
        ret['category_type'] = category_type
        ret['category_label'] = category_label
        ret['landmark_vis'] = landmark_vis
        ret['landmark_pos'] = landmark_pos
        ret['landmark_pos_normalized'] = landmark_pos_normalized
        ret['attr'] = attr

        return ret

        # return image, category_type, category_label, landmark_vis, landmark_pos, attr

    def __len__(self):
        return len(self.df)
